# coding=utf-8
import os
import json
import pickle


class ParseInfo:
    # code 证券代码 name 证券名称 price 最近成交价 volume 成交量 amount 成交额
    def __init__(
        self,
        code: str,
        name: str,
        price: float,
        volume: float,
        amount: float,
        last_close: float,  # yesterday close
    ):
        self.code = code
        self.name = name
        self.price = price
        self.vol = volume
        self.amount = amount
        self.ystd_close = last_close

    def output(self):
        print(json.dumps({
            'code': self.code,
            'name': self.name,
            'price': self.price,
            'volume': self.vol,
            'amount': self.amount,
        }, ensure_ascii=False))


class KLine:
    def __init__(
        self,
        name: str = '',
        p_open: float = None,
        p_close: float = None,
        p_high: float = None,
        p_low: float = None,
        prev_close: float = None,
        prev_amount: float = None,
        prev_vol: float = None,
        last_close: float = None,  # yesterday close
        amount: float = float(0.0),
        vol: int = int(0),
    ):
        self.name = name

        self.open = p_open
        self.close = p_close
        self.high = p_high
        self.low = p_low

        self.amount = amount
        self.vol = vol

        self.prev_close = prev_close
        self.prev_amount = prev_amount
        self.prev_vol = prev_vol

        self.ystd_close = last_close

    def update_prices(self, price: float) -> None:
        if self.open is None:
            self.open = price

        self.close = price

        if self.high is None:
            self.high = price
        else:
            self.high = max(self.high, price)

        if self.low is None:
            self.low = price
        else:
            self.low = min(self.low, price)

    def prepare_output_kline(self):
        new_kline = KLine(name=self.name)

        if (self.open is None or self.open < 0.001) \
                and (self.close is None or self.close < 0.001) \
                and (self.high is None or self.high < 0.001) \
                and (self.low is None or self.low < 0.001):

            t_prev_close = self.prev_close
            if t_prev_close is None:
                t_prev_close = self.ystd_close

            new_kline.open = t_prev_close
            new_kline.close = t_prev_close
            new_kline.high = t_prev_close
            new_kline.low = t_prev_close
        else:
            # open
            if self.open is None or self.open < 0.001:
                new_kline.open = self.close
            else:
                new_kline.open = self.open

            # low
            if self.low is None or self.low < 0.001:
                new_kline.low = self.close
            else:
                new_kline.low = self.low

            # high
            new_kline.high = self.high

            # close
            new_kline.close = self.close

        return new_kline

    def output(self):
        print(self.name, self.high, self.low, self.open, self.close, self.amount, self.vol)


class KLineList:
    def __init__(self):
        self.klines: dict[str, KLine] = {}  # {code: kline}

    def initialize_by_previous_kline_list(self, prev_klines_list):
        if prev_klines_list is not None:
            for key in prev_klines_list.klines.keys():
                self.klines[key] = KLine(name=prev_klines_list.klines[key].name)
                self.klines[key].prev_close = prev_klines_list.klines[key].close

    def update_kline_by_parse_infos(self, info_list: list):
        for info in info_list:
            if info.code not in self.klines.keys():
                self.klines[info.code] = KLine()
                self.klines[info.code].name = info.name
            self.klines[info.code].update_prices(info.price)
            self.klines[info.code].amount = info.amount
            self.klines[info.code].vol = info.vol
            self.klines[info.code].ystd_close = info.ystd_close

    def output(self):
        print('============')
        for k in self.klines.keys():
            kline = self.klines[k]
            print(k, kline.name,
                  'high: ' + str(kline.high),
                  'low: ' + str(kline.low),
                  'open: ' + str(kline.open),
                  'close: ' + str(kline.close),
                  'amount: ' + str(kline.amount),
                  'vol: ' + str(kline.vol),
                  'prev_close: ' + str(kline.prev_close))


def save_backup(backup_path: str, klines):
    klines_backup = pickle.dumps(klines)
    with open(backup_path, 'w') as f:
        f.write(klines_backup)


def load_backup(backup_path: str):
    if os.path.exists(backup_path):
        with open(backup_path, 'r') as f:
            b = f.read()
            return pickle.loads(b)
    else:
        print('Did not find the backup path: ' + backup_path)
        return None


def test_kline_list():
    BACKUP_PATH = './_cache/backup.bak'

    klines = KLineList()

    klines.update_kline_by_parse_infos([
        ParseInfo('sh000001', '上证指数', price=12.0, volume=10, amount=1000, last_close=9.0),
        ParseInfo('sz000001', '深证指数', price=20.0, volume=20, amount=2000, last_close=9.0),
    ])
    # klines.output()

    klines.update_kline_by_parse_infos([
        ParseInfo('sh000001', '上证指数', price=14.0, volume=11, amount=1001, last_close=9.0),
        ParseInfo('sz000001', '深证指数', price=19.0, volume=21, amount=2001, last_close=9.0),
    ])
    # klines.output()

    klines.update_kline_by_parse_infos([
        ParseInfo('sh000001', '上证指数', price=10.0, volume=12, amount=1002, last_close=9.0),
        ParseInfo('sz000001', '深证指数', price=21.0, volume=22, amount=2002, last_close=9.0),
    ])
    klines.output()

    save_backup(BACKUP_PATH, klines)

    c = load_backup(BACKUP_PATH)
    c.output()

    klines_new = KLineList()
    klines_new.initialize_by_previous_kline_list(klines)
    klines_new.output()


if __name__ == '__main__':
    test_kline_list()
